import numpy as np
from typing import List, Tuple


def tsp_exact_dp(dist_matrix: np.ndarray) -> Tuple[float, List[int]]:
    """
    使用动态规划求解TSP的最优解

    Args:
        dist_matrix: 距离矩阵，shape=(n, n)，对角线为inf

    Returns:
        (最优路径长度, 最优路径节点序列)
    """
    n = dist_matrix.shape[0]

    dp = np.full((1 << n, n), float('inf'))

    for i in range(1, n):
        dp[1 << i][i] = dist_matrix[0][i]

    for mask in range(1 << n):
        for u in range(n):
            if not (mask & (1 << u)):
                continue
            if dp[mask][u] == float('inf'):
                continue

            for v in range(n):
                if mask & (1 << v):
                    continue

                new_mask = mask | (1 << v)
                new_cost = dp[mask][u] + dist_matrix[u][v]
                if new_cost < dp[new_mask][v]:
                    dp[new_mask][v] = new_cost

    final_mask = (1 << n) - 1
    best_length = float('inf')
    best_end_city = -1

    for i in range(1, n):
        cost = dp[final_mask][i] + dist_matrix[i][0]
        if cost < best_length:
            best_length = cost
            best_end_city = i

    path = [0]
    mask = final_mask
    current = best_end_city

    while current != 0:
        path.append(current)

        prev_mask = mask ^ (1 << current)
        prev_city = -1
        min_cost = float('inf')

        for u in range(n):
            if not (prev_mask & (1 << u)):
                continue
            cost = dp[prev_mask][u] + dist_matrix[u][current]
            if cost < min_cost:
                min_cost = cost
                prev_city = u

        current = prev_city
        mask = prev_mask

    return best_length, path[::-1]


def read_tsp_file(filepath: str) -> np.ndarray:
    """
    读取标准TSP格式文件，返回距离矩阵
    """
    with open(filepath, 'r') as f:
        lines = f.readlines()

    coord_start = None
    for i, line in enumerate(lines):
        if "NODE_COORD_SECTION" in line:
            coord_start = i + 1
            break

    coords = []
    for i in range(coord_start, len(lines)):
        line = lines[i].strip()
        if line == "EOF":
            break
        parts = line.split()
        if len(parts) >= 3:
            node_id = int(parts[0])
            x = float(parts[1])
            y = float(parts[2])
            coords.append([x, y])

    coords = np.array(coords)

    # 计算距离矩阵
    n = len(coords)
    dist_matrix = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            if i == j:
                dist_matrix[i][j] = float('inf')
            else:
                dist_matrix[i][j] = np.sqrt(np.sum((coords[i] - coords[j]) ** 2))

    return dist_matrix


if __name__ == "__main__":
    filepath = "../data/random/random_tsp_10_1762100645.tsp"

    print(f"正在读取TSP文件: {filepath}")
    dist_matrix = read_tsp_file(filepath)

    print(f"城市数量: {dist_matrix.shape[0]}")

    best_length, best_path = tsp_exact_dp(dist_matrix)

    print("基于精确算法的结果")
    print(f"最优路径长度: {best_length:.2f}")
    print(f"最优路径: {' -> '.join(map(str, best_path))}")
